package com.hulk.rpc.serialize.impl.kryo;

import com.hulk.rpc.config.common.HulkConstants;
import com.hulk.rpc.param.MethodParam;
import com.hulk.rpc.protocol.Request;
import com.hulk.rpc.protocol.Response;
import com.hulk.rpc.protocol.ResponseStatus;
import com.hulk.rpc.protocol.recycle.RecycleRequest;
import com.hulk.rpc.protocol.recycle.RecycleResponse;
import com.hulk.rpc.recycle.RecycleUtils;
import com.hulk.rpc.serialize.Serializer;
import com.hulk.rpc.serialize.TracerSerializer;
import com.hulk.rpc.trace.Tracer;
import com.hulk.rpc.utils.ByteBufUtils;
import io.netty.buffer.ByteBuf;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.Map;
import java.util.function.Supplier;

@Slf4j
public class KryoSerializer extends Serializer {
  private Map<Class<?>, Integer> registerMap;

  private final Supplier<KryoContext> producer = KryoContext::new;

  private final ThreadLocal<KryoContext> kryoContextHolder = ThreadLocal.withInitial(producer);
  private final TracerSerializer tracerSerializer = new TracerSerializer();

  @Override
  public boolean isSupportedClassId() {
    return true;
  }

  @Override
  public void setClassIds(Map<Class<?>, Integer> classIds) {
    this.registerMap = classIds;
  }

  private KryoContext kryoContext() {
    KryoContext kryoContext = kryoContextHolder.get();
    kryoContext.registerClassIds(registerMap);
    return kryoContext;
  }

  @Override
  public void writeRequest(ByteBuf byteBuf, Request request) throws IOException {
    final int beginWriterIndex = byteBuf.writerIndex();

    byteBuf.writeInt(0);
    // byteBuf.writerIndex(beginWriterIndex + HulkConstants.HEADER_FIELD_LENGTH);
    byteBuf.writeInt(request.getRequestId());
    ByteBufUtils.writeVarInt(byteBuf, request.getServiceId());
    tracerSerializer.write(byteBuf, request.getTracer());

    if (request.getMethodParam() == null) {
      byteBuf.writeBoolean(false);
    } else {
      byteBuf.writeBoolean(true);
      kryoContext().writeObject(byteBuf, request.getMethodParam());
    }

    int finishWriterIndex = byteBuf.writerIndex();
    int length = finishWriterIndex - beginWriterIndex - HulkConstants.HEADER_FIELD_LENGTH;

    byteBuf.setInt(beginWriterIndex, length);

    RecycleUtils.release(request);
  }

  @Override
  public Request readRequest(ByteBuf byteBuf) throws IOException {
    int requestId = byteBuf.readInt();
    int serviceId = ByteBufUtils.readVarInt(byteBuf);
    Tracer tracer = tracerSerializer.read(byteBuf);

    MethodParam methodParam = null;
    if (byteBuf.readBoolean()) {
      Class<?> clazz = getMethodParamClass(serviceId);
      methodParam = (MethodParam) kryoContext().readObject(byteBuf, clazz);
    }

    return RecycleRequest.newInstance(requestId, serviceId, tracer, methodParam);
  }

  @Override
  public void writeResponse(ByteBuf byteBuf, Response response) throws IOException {
    final int beginWriterIndex = byteBuf.writerIndex();

    byteBuf.writeInt(0);
    // byteBuf.writerIndex(beginWriterIndex + HulkConstants.HEADER_FIELD_LENGTH);
    byteBuf.writeInt(response.getRequestId());

    final int statusWriterIndex = byteBuf.writerIndex();

    try {
      byteBuf.writeByte(response.getStatusCode());
      tracerSerializer.write(byteBuf, response.getTracer());
      kryoContext().writeClassAndObject(byteBuf, response.getResult());
    } catch (Exception e) {
      if (log.isWarnEnabled()) {
        log.warn("kryo writeResponse error", e);
      }

      byteBuf.writerIndex(statusWriterIndex);
      byteBuf.writeByte(ResponseStatus.BAD_RESPONSE);
      tracerSerializer.write(byteBuf, response.getTracer());
      kryoContext().writeClassAndObject(byteBuf, e.getMessage());
    }

    int finishWriterIndex = byteBuf.writerIndex();
    int length = finishWriterIndex - beginWriterIndex - HulkConstants.HEADER_FIELD_LENGTH;

    byteBuf.setInt(beginWriterIndex, length);

    RecycleUtils.release(response);
  }

  @Override
  public Response readResponse(ByteBuf byteBuf) throws IOException {
    int requestId = byteBuf.readInt();
    byte statusCode = byteBuf.readByte();
    Tracer tracer = tracerSerializer.read(byteBuf);

    Object result = kryoContext().readClassAndObject(byteBuf);

    return RecycleResponse.newInstance(requestId, statusCode, tracer, result);
  }

}
